from Crypto.Cipher import AES

from Crypto.Util.Padding import pad, unpad
from typing import Tuple
import json
import base64
from typing import Dict, List, Tuple
import os

from scorpio.core.common.logger import get_logger
logger = get_logger(__name__)

LOCAL_KEY_CACHE = [-70, -69, 74, -97, 119, 74, -72, 83, -55, 108, 45, 101, 61, -2, 84, 74]
default_key = bytes((b + 256) % 256 for b in LOCAL_KEY_CACHE)
default_workspace_dir=f"{os.getenv("HOME", "")}/.cloudbeaver/workspace"

class DefaultValueDecryptor:
    key: bytes = default_key
    workspace_dir: str = default_workspace_dir 
    dbeaver_path: str
    username: str
    filename: str
    data: bytes = None
    def __init__(self, username: str, credential: str = "credentials-config.json", dbeaver_path: str = None):
        self.username = username
        self.dbeaver_path = dbeaver_path
        if dbeaver_path is not None:
            self.credential_path = f"{dbeaver_path}/{self.username}/.dbeaver/{credential}"
        # set none if not find the file 
    def _read(self, credential_path: str)->bytes:
        with open(credential_path, 'rb') as f:
            data = f.read()
        return data
    def decrypt(self, data: bytes):
        if data is None: 
            data = self._read(self.credential_path)
        else:
            pass
        if len(data)>=16:
            ok, result=self._decrypt_cbc(data)
        else:
            if len(data) > 4:
                # 尝试跳过前4字节
                ok, result=self._decrypt_cbc(data)
            else:
                pass

        if ok:
            return ok, result
        else:
            try:
                cipher = AES.new(self.key, AES.MODE_ECB)
                decrypted_data = cipher.decrypt(data)
                # 尝试去除填充
                try:
                    result = unpad(decrypted_data, AES.block_size)
                except:
                    result = decrypted_data
                return True, result.decode("utf-8")
            except Exception as e:
                return False, f"failed to decrypt under ECB model: {str(e)}"
    # CBC模式解密辅助函数
    def _decrypt_cbc(self, data)->Tuple[bool, str]:
        iv = data[:16]
        cipher_text = data[16:]
        logger.debug(f"IV: {iv.hex()}, 密文长度: {len(cipher_text)}")
        # 检查长度
        # 密文长度不是16的倍数, 需要添加填充
        if len(cipher_text) % 16 != 0:
            # 尝试修复：添加最小填充
            padding_needed = 16 - (len(cipher_text) % 16)
            cipher_text += bytes([0] * padding_needed)
        try:
            cipher = AES.new(self.key, AES.MODE_CBC, iv=iv)
            decrypted_data = cipher.decrypt(cipher_text)
            try:
                result = unpad(decrypted_data, AES.block_size)
            except ValueError as e:
                logger.warning(f"failed to unpad, error: {e}")
                result = decrypted_data
            return True, result.decode("utf-8")
        except Exception as e:
            return False, f"failed to decrypt under CBC model, error: {str(e)}"

class IngestDBConfig:
    key: bytes = default_key
    workspace_dir: str 
    def __init__(self, fingerprint: str) -> None:
        self.home=os.getenv("HOME", ".")
        self.workspace_dir=os.getenv("CLOUDBEAVER_WORKSPACE_DIR", f"{self.home}/.cloudbeaver/workspace")
        array=fingerprint.split(" ")
        self.bearer=True
        if array[0]=="Bearer":
            self.fingerprint=array[1]
        else:
            self.bearer=False
            self.fingerprint=array[0]
    def decode_fingerprint(self)->Dict:
        # 将 Base64 字符串解码为字节
        body=self.fingerprint
        db_instance_str = base64.b64decode(body).decode('utf-8')
        
        # 将 JSON 字符串解析为字典
        db_instance: Dict= json.loads(db_instance_str)
        return db_instance

    def get_config(self)->Tuple[str, Dict]:
        db_instance: Dict=self.decode_fingerprint()
        username: str=db_instance["username"]
        connection_uuid: str = db_instance["#connection"]
        dbeaver_path=f"{self.workspace_dir}/user-projects/{username}/.dbeaver"
        data_sources_path=f"{dbeaver_path}/data-sources.json"
        with open(data_sources_path) as f:
            data_sources_data: Dict= json.load(f)
        #print(data_sources_data["connections"])
        connection: Dict = None
        _conn=data_sources_data["connections"].get(connection_uuid)
        if _conn:
            connection=_conn["configuration"]
        if connection is not None:
            decryptor=DefaultValueDecryptor(username=username, dbeaver_path=dbeaver_path)
            credential=self.get_bear_credential(decryptor, dbeaver_path, connection_uuid)
            if credential is not None:
                connection["user"]=credential["user"]
                #print(f"{connection}\n,{credential}")
                if _conn.get("save-password"):
                    connection["password"]=credential["password"]
                else:
                    pass
            else:
                pass
        return connection_uuid, connection
    def get_bear_config(self)->Tuple[str, Dict]:
        db_instance: Dict=self.decode_fingerprint()
        username: str=db_instance["username"]
        dbeaver_path=f"{self.workspace_dir}/user-projects/{username}/.dbeaver"
        data_sources_path=f"{dbeaver_path}/data-sources.json"
        with open(data_sources_path) as f:
            data_sources_data: Dict= json.load(f)
        #print(data_sources_data["connections"])
        connection=db_instance["connection"]
        _host = connection["host"]
        _port = int(connection["port"])  # 确保为整数
        _db = connection.get("database")  # 可能为None
        connection_uuid: str = None
        for key, _conn in data_sources_data["connections"].items():
            config=_conn['configuration']
            host = config.get("host")
            port = int(config.get("port", "0"))  # 转换为整数
            db = config.get("database")
            
            host_match = host == _host
            port_match = port == _port
            db_match = True
            if _db is not None and db is not None:
                db_match = db == _db
            if host_match and port_match and db_match:
                connection_uuid = key
                connection = config
                break
        #print(f"\n###\ncapture configure:\n{connection_uuid}\n{connection}\n\n")
        if connection_uuid is not None:
            decryptor=DefaultValueDecryptor(username=username, dbeaver_path=dbeaver_path)
            credential=self.get_bear_credential(decryptor, dbeaver_path, connection_uuid)
            if credential is not None:
                connection["user"]=credential["user"]
                if credential.get("password"):
                    connection["password"]=credential["password"]
                else:
                    pass
            else:
                pass
        return connection_uuid, connection
    def get_bear_credential(self, decryptor: DefaultValueDecryptor, dbeaver_path: str, connection_uuid: str)->Tuple[str, Dict]:
        credentials_config=f"{dbeaver_path}/credentials-config.json"
        with open(credentials_config, 'rb') as f:
            data: bytes= f.read()
        ok, credentials=decryptor.decrypt(data=data)
        if not ok:
            return None
        result: Dict=json.loads(credentials)
        credential: Dict=result.get(connection_uuid)
        return credential['#connection']

if __name__ == "__main__":
    ingestion=IngestDBConfig("Bearer eyJ1c2VybmFtZSI6ICJyb290IiwgImNvbm5lY3Rpb24iOiB7Imhvc3QiOiAiMTcyLjExLjYuMjEiLCAicG9ydCI6IDMzMDZ9fQ==") 
    connection_uuid,db_config=ingestion.get_bear_config()
    print(f"ingest Bear db_config from cloudbeaver: {db_config}")
    db_instance={
        "username": 'root',
        "#connection": "postgres-jdbc-197e97b8150-263272c8d558544b"
    }
    
    # 将 JSON 字符串编码为 Base64
    db_instance_bytes = json.dumps(db_instance).encode('utf-8')
    fingerprint = base64.b64encode(db_instance_bytes).decode('utf-8')
    print(fingerprint)
    ingestion=IngestDBConfig(fingerprint)
    connection_uuid, db_config=ingestion.get_config()
    print(f"ingest db_config from cloudbeaver: {db_config}")
